import torch
import copy
import torch.nn as nn
from torch.nn import functional as F
from torch.nn import Module
from torch.nn import MultiheadAttention  # 多头注意力机制
from torch.nn import ModuleList  # 用于存储多个模块的列表
from torch.nn.init import xavier_uniform_  # 参数初始化方法
from torch.nn import Dropout  # dropout层，防止过拟合
from torch.nn import Linear  # 线性层
from torch.nn import LayerNorm  # 层归一化


class Encoder(Module):
    r"""编码器，由N个编码器层堆叠而成

    Args:
        encoder_layer: EncoderLayer类的实例（必需），单个编码器层
        num_layers: 编码器中层的数量（必需）
        norm: 层归一化组件（可选）
    """

    def __init__(self, encoder_layer, num_layers, norm=None):
        super(Encoder, self).__init__()
        # 复制num_layers个编码器层，组成编码器堆叠
        self.layers = _get_clones(encoder_layer, num_layers)
        self.num_layers = num_layers  # 层数
        self.norm = norm  # 归一化层

    def forward(self, src):
        r"""依次通过所有编码器层处理输入

        Args:
            src: 输入特征张量
        Returns:
            经过所有编码器层处理后的输出
        """
        output = src  # 初始输出为输入

        # 逐层处理
        for i in range(self.num_layers):
            output = self.layers[i](output)

        # 若指定了归一化层，对最终输出进行归一化
        if self.norm:
            output = self.norm(output)

        return output


class Decoder(Module):
    r"""解码器，由N个解码器层堆叠而成

    Args:
        decoder_layer: DecoderLayer类的实例（必需），单个解码器层
        num_layers: 解码器中层的数量（必需）
        norm: 层归一化组件（可选）
    """

    def __init__(self, decoder_layer, num_layers, norm=None):
        super(Decoder, self).__init__()
        # 复制num_layers个解码器层，组成解码器堆叠
        self.layers = _get_clones(decoder_layer, num_layers)
        self.num_layers = num_layers  # 层数
        self.norm = norm  # 归一化层

    def forward(self, tgt, memory):
        r"""依次通过所有解码器层处理输入和记忆特征

        Args:
            tgt: 目标序列特征（解码器输入）
            memory: 编码器输出的记忆特征
        Returns:
            经过所有解码器层处理后的输出
        """
        output = tgt  # 初始输出为目标输入

        # 逐层处理
        for i in range(self.num_layers):
            output = self.layers[i](output, memory)

        # 若指定了归一化层，对最终输出进行归一化
        if self.norm:
            output = self.norm(output)

        return output


class EncoderLayer(Module):
    r"""编码器层，借鉴自CMRAN模型

    Args:
        d_model: 输入特征的维度（必需）
        nhead: 多头注意力机制的头数（必需）
        dim_feedforward: 前馈网络的隐藏层维度（默认=1024）
        dropout: dropout概率（默认=0.1）
        activation: 中间层的激活函数，可选relu或gelu（默认=relu）
    """

    def __init__(self, d_model, nhead, dim_feedforward=1024, dropout=0.1, activation="relu"):
        super(EncoderLayer, self).__init__()
        # 自注意力机制（多头）
        self.self_attn = MultiheadAttention(d_model, nhead, dropout=dropout)
        # 前馈网络实现
        self.linear1 = Linear(d_model, dim_feedforward)  # 升维线性层
        self.dropout = Dropout(dropout)  # dropout层
        self.linear2 = Linear(dim_feedforward, d_model)  # 降维线性层

        # 归一化层
        self.norm1 = LayerNorm(d_model)
        self.norm2 = LayerNorm(d_model)
        # 注意力和前馈网络后的dropout
        self.dropout1 = Dropout(dropout)
        self.dropout2 = Dropout(dropout)

        # 获取激活函数
        self.activation = _get_activation_fn(activation)

    def forward(self, src):
        r"""通过编码器层处理输入

        Args:
            src: 输入特征张量
        Returns:
            经过编码器层处理后的输出
        """
        # 自注意力计算：query=key=value=src
        src2 = self.self_attn(src, src, src)[0]
        # 残差连接 + dropout
        src = src + self.dropout1(src2)
        # 层归一化
        src = self.norm1(src)

        # 前馈网络
        if hasattr(self, "activation"):
            src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))
        else:  # 向后兼容（若未指定激活函数则默认用relu）
            src2 = self.linear2(self.dropout(F.relu(self.linear1(src))))
        # 残差连接 + dropout
        src = src + self.dropout2(src2)
        # 层归一化
        src = self.norm2(src)
        return src


class DecoderLayer(Module):
    r"""解码器层，借鉴自CMRAN模型

    Args:
        d_model: 输入特征的维度（必需）
        nhead: 多头注意力机制的头数（必需）
        dim_feedforward: 前馈网络的隐藏层维度（默认=1024）
        dropout: dropout概率（默认=0.1）
        activation: 中间层的激活函数，可选relu或gelu（默认=relu）
    """

    def __init__(self, d_model, nhead, dim_feedforward=1024, dropout=0.1, activation="relu"):
        super(DecoderLayer, self).__init__()
        # 自注意力机制（目标序列内部）
        self.self_attn = MultiheadAttention(d_model, nhead, dropout=dropout)
        # 多头注意力机制（目标序列与记忆特征之间）
        self.multihead_attn = MultiheadAttention(d_model, nhead, dropout=dropout)
        # 前馈网络实现
        self.linear1 = Linear(d_model, dim_feedforward)  # 升维线性层
        self.dropout = Dropout(dropout)  # dropout层
        self.linear2 = Linear(dim_feedforward, d_model)  # 降维线性层

        # 归一化层
        self.norm1 = LayerNorm(d_model)
        self.norm2 = LayerNorm(d_model)
        # 注意力和前馈网络后的dropout
        self.dropout1 = Dropout(dropout)
        self.dropout2 = Dropout(dropout)

        # 获取激活函数
        self.activation = _get_activation_fn(activation)

    def forward(self, tgt, memory):
        r"""通过解码器层处理输入和记忆特征

        Args:
            tgt: 目标序列特征（解码器输入）
            memory: 编码器输出的记忆特征
        Returns:
            经过解码器层处理后的输出
        """
        # 将记忆特征与目标特征拼接（作为注意力的key和value）
        memory = torch.cat([memory, tgt], dim=0)
        # 计算目标序列与记忆特征的注意力
        tgt2 = self.multihead_attn(tgt, memory, memory)[0]
        # 残差连接 + dropout
        tgt = tgt + self.dropout1(tgt2)
        # 层归一化
        tgt = self.norm1(tgt)

        # 前馈网络
        if hasattr(self, "activation"):
            tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
        else:  # 向后兼容
            tgt2 = self.linear2(self.dropout(F.relu(self.linear1(tgt))))
        # 残差连接 + dropout
        tgt = tgt + self.dropout2(tgt2)
        # 层归一化
        tgt = self.norm2(tgt)
        return tgt


def _get_clones(module, N):
    r"""复制N个相同的模块，组成ModuleList

    Args:
        module: 要复制的模块
        N: 复制的数量
    Returns:
        包含N个复制模块的ModuleList
    """
    return ModuleList([copy.deepcopy(module) for i in range(N)])


def _get_activation_fn(activation):
    r"""根据名称获取激活函数

    Args:
        activation: 激活函数名称（"relu"或"gelu"）
    Returns:
        对应的激活函数
    Raises:
        RuntimeError: 若激活函数名称不支持
    """
    if activation == "relu":
        return F.relu
    elif activation == "gelu":
        return F.gelu
    else:
        raise RuntimeError("activation should be relu/gelu, not %s." % activation)


class New_Audio_Guided_Attention(nn.Module):
    r"""
    新的音频引导注意力机制，与论文中描述略有不同，后续发现此实现更高效。
    功能：结合音频特征引导视频特征的注意力计算，增强跨模态特征交互
    """
   
    def __init__(self, beta):
        super(New_Audio_Guided_Attention, self).__init__()

        self.beta = beta  # 自注意力特征的权重系数
        self.relu = nn.ReLU()  # ReLU激活函数
        self.video_input_dim = 512  # 视频输入特征维度
        self.audio_input_dim = 128  # 音频输入特征维度

        self.hidden_dim = 256  # 隐藏层维度
        # 通道注意力相关层
        self.affine_video_1 = nn.Linear(self.video_input_dim, self.video_input_dim)  # 视频特征线性变换
        self.affine_audio_1 = nn.Linear(self.audio_input_dim, self.video_input_dim)  # 音频特征线性变换（映射到视频维度）
        self.affine_bottleneck = nn.Linear(self.video_input_dim, self.hidden_dim)  # 瓶颈层（降维）
        self.affine_v_c_att = nn.Linear(self.hidden_dim, self.video_input_dim)  # 输出通道注意力权重

        # 空间注意力相关层
        self.affine_video_2 = nn.Linear(self.video_input_dim, self.hidden_dim)  # 视频特征线性变换
        self.affine_audio_2 = nn.Linear(self.audio_input_dim, self.hidden_dim)  # 音频特征线性变换（映射到隐藏层维度）
        self.affine_v_s_att = nn.Linear(self.hidden_dim, 1)  # 输出空间注意力权重

        # 自注意力相关层
        self.latent_dim = 4  #  latent维度（用于自注意力降维）
        self.video_query = nn.Linear(self.video_input_dim, self.video_input_dim//self.latent_dim)  # 视频query映射
        self.video_key = nn.Linear(self.video_input_dim, self.video_input_dim//self.latent_dim)   # 视频key映射
        self.video_value = nn.Linear(self.video_input_dim, self.video_input_dim)  # 视频value映射

        # 视频平均特征相关层（用于空间注意力）
        self.affine_video_ave = nn.Linear(self.video_input_dim, self.hidden_dim)  # 平均视频特征映射
        self.affine_video_3 = nn.Linear(self.video_input_dim, self.hidden_dim)    # 视频特征映射
        self.ave_bottleneck = nn.Linear(512, 256)  # 平均特征瓶颈层（未在forward中使用，可能为预留）
        self.ave_v_att = nn.Linear(self.hidden_dim, 1)  # 输出平均特征引导的注意力权重

        self.tanh = nn.Tanh()  # Tanh激活函数
        self.softmax = nn.Softmax(dim=-1)  # 用于计算注意力权重的softmax
        self.dropout = nn.Dropout(0.2)  # dropout层（概率0.2）
        self.norm = nn.LayerNorm(self.video_input_dim)  # 视频特征归一化层


    def forward(self, video, audio):
        '''
        前向传播：结合音频特征计算视频特征的注意力，输出增强后的视频特征

        :param video: 视频特征，形状为 [batch, 10, 7, 7, 512] 
                      （batch：批次大小，10：时间步，7x7：空间维度，512：特征维度）
        :param audio: 音频特征，形状为 [batch, 10, 128] 
                      （batch：批次大小，10：时间步，128：特征维度）
        :return: 增强后的视频特征，形状为 [batch, 10, 512]
        '''
        # 调整音频特征维度：[batch, 10, 128] -> [10, batch, 128]（可能为适配后续操作）
        audio = audio.transpose(1, 0)
        # 获取视频特征各维度大小
        batch, t_size, h, w, v_dim = video.size()  # t_size=10, h=7, w=7, v_dim=512
        a_dim = audio.size(-1)  # 音频特征维度=128
        # 音频特征重塑：[10, batch, 128] -> [batch*10, 128]（合并批次和时间步）
        audio_feature = audio.reshape(batch * t_size, a_dim)
        # 视频特征重塑：[batch, 10, 7, 7, 512] -> [batch, 10, 49, 512]（7x7空间维度展平为49）
        visual_feature = video.reshape(batch, t_size, -1, v_dim)
        raw_visual_feature = visual_feature  # 保存原始视频特征用于后续残差

        # ============================== 自注意力计算 =======================================
        # 视频query特征：[batch, 10, 49, 512] -> [batch*10, 49, 128]（512/4=128）
        video_query_feature = self.video_query(visual_feature).reshape(batch * t_size, h * w, -1)
        # 视频key特征：[batch, 10, 49, 512] -> [batch*10, 49, 128] -> 转置为[batch*10, 128, 49]
        video_key_feature = self.video_key(visual_feature).reshape(batch * t_size, h * w, -1).permute(0, 2, 1)
        # 计算注意力分数：query与key矩阵乘法 -> [batch*10, 49, 49]
        energy = torch.bmm(video_query_feature, video_key_feature)
        # 归一化得到注意力权重
        attention = self.softmax(energy)
        # 视频value特征：[batch, 10, 49, 512] -> [batch*10, 49, 512]
        video_value_feature = self.video_value(visual_feature).reshape(batch * t_size, h * w, -1)
        # 注意力加权求和：[batch*10, 49, 49] * [batch*10, 49, 512] -> [batch*10, 49, 512]
        output = torch.matmul(attention, video_value_feature)
        # 残差连接 + dropout + 层归一化
        output = self.norm(visual_feature.reshape(batch * t_size, h * w, -1) + self.dropout(output))
        visual_feature = output  # 更新视频特征为自注意力处理后的值

        # ============================== 视频自空间注意力 ====================================
        # 计算视频特征的空间平均（对49个空间位置求和后平均）-> [batch*10, 512]
        video_average = visual_feature.sum(dim=1)/(h*w)
        video_average = video_average.reshape(batch*t_size, v_dim)
        # 平均特征映射并增加维度 -> [batch*10, 1, 256]
        video_average = self.relu(self.affine_video_ave(video_average)).unsqueeze(-2)
        # 视频特征重塑 -> [batch*10, 49, 512]
        self_video_att_feat = visual_feature.reshape(batch*t_size, -1, v_dim)
        # 视频特征映射到隐藏层 -> [batch*10, 49, 256]
        self_video_att_query = self.relu(self.affine_video_3(self_video_att_feat))
        # 平均特征与视频特征交互（元素乘法）-> [batch*10, 49, 256]
        self_query = self_video_att_query * video_average
        # 计算空间注意力权重并归一化 -> [batch*10, 1, 49]
        self_spatial_att_maps = self.softmax(self.tanh(self.ave_v_att(self_query)).transpose(2,1))
        # 注意力加权求和得到自空间注意力特征 -> [batch, 10, 512]
        self_att_feat = torch.bmm(self_spatial_att_maps, visual_feature).squeeze().reshape(batch, t_size, v_dim)


        # ============================== 通道注意力 ====================================
        # 音频特征映射并增加维度 -> [batch*10, 1, 512]
        audio_query_1 = self.relu(self.affine_audio_1(audio_feature)).unsqueeze(-2)
        # 视频特征映射 -> [batch*10, 49, 512]
        video_query_1 = self.relu(self.affine_video_1(visual_feature)).reshape(batch*t_size, h*w, -1)
        # 音频与视频特征交互（元素乘法）并在空间维度平均 -> [batch*10, 512]
        audio_video_query_raw = (audio_query_1 * video_query_1).mean(-2)
        # 瓶颈层降维 -> [batch*10, 256]
        audio_video_query = self.relu(self.affine_bottleneck(audio_video_query_raw))
        # 计算通道注意力权重（sigmoid确保在0-1之间）-> [batch, 10, 49, 512]
        channel_att_maps = self.affine_v_c_att(audio_video_query).sigmoid().reshape(batch, t_size, -1, v_dim)
        # 应用通道注意力（原始视频特征 * (注意力权重 + 1) 实现残差）
        c_att_visual_feat = (raw_visual_feature * (channel_att_maps + 1))


        # ============================== 音频引导空间注意力 =====================================
        # 通道注意力处理后的视频特征重塑 -> [batch*10, 49, 512]
        c_att_visual_feat = c_att_visual_feat.reshape(batch*t_size, -1, v_dim)
        # 视频特征映射到隐藏层 -> [batch*10, 49, 256]
        c_att_visual_query = self.relu(self.affine_video_2(c_att_visual_feat))
        # 音频特征映射并增加维度 -> [batch*10, 1, 256]
        audio_query_2 = self.relu(self.affine_audio_2(audio_feature)).unsqueeze(-2)
        # 音频与视频特征交互（元素乘法）-> [batch*10, 49, 256]
        audio_video_query_2 = c_att_visual_query * audio_query_2
        # 计算空间注意力权重并归一化 -> [batch*10, 1, 49]
        spatial_att_maps = self.softmax(self.tanh(self.affine_v_s_att(audio_video_query_2)).transpose(2, 1))
        # 注意力加权求和得到音频引导空间注意力特征 -> [batch, 10, 512]
        c_s_att_visual_feat = torch.bmm(spatial_att_maps, c_att_visual_feat).squeeze().reshape(batch, t_size, v_dim)

        # 结合自注意力特征和音频引导特征（残差连接）
        c_s_att_visual_feat = c_s_att_visual_feat + self.beta * self_att_feat.sigmoid() * c_s_att_visual_feat


        return c_s_att_visual_feat